import json
import os

import pandas as pd
import qdrant_client
from openai import OpenAI
from qdrant_client.http import models as rest

api_key = ''
# 实例化客户端
client = OpenAI(api_key=api_key,
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
GPT_MODEL = "qwen-max"
EMBEDDING_MODEL = "text-embedding-v2"

records = []

for x in records_list:
    records_path = "D:/robotlearn/agent/Customs_agent_assistant/data/虚拟报关记录/" + x

    # Opening JSON file
    f = open(records_path, encoding='utf-8')

    # returns JSON object as
    # a dictionary
    data = json.load(f)

    records.append(data)

    # Closing file
    f.close()

for i, record in enumerate(records):
    # 拼接除货物数量和货物价值之外的所有变量作为文本
    text_to_embed = f"报关单号：{record['报关单号']}，发货人：{record['发货人']}，发货地：{record['发货地']}，收货地：{record['收货地']}，货物名称：{record['货物名称']}，运输方式：{record['运输方式']}，报关日期：{record['报关日期']}"
    try:
        embedding = client.embeddings.create(model=EMBEDDING_MODEL, input=text_to_embed)
        records[i].update({"embedding": embedding.data[0].embedding})
    except Exception as e:
        print(f"处理记录 {record['报关单号']} 时出错：{e}")

qdrant = qdrant_client.QdrantClient(host="localhost", port=6333)
qdrant.get_collections()

collection_name = "Customs declaration records"

vector_size = len(records[0]["embedding"])
vector_size

article_df = pd.DataFrame(records)
article_df.head()

# Delete the collection if it exists, so we can rewrite it changes to articles were made
#if qdrant.get_collection(collection_name=collection_name):
    #qdrant.delete_collection(collection_name=collection_name)

# Create Vector DB collection
qdrant.create_collection(
    collection_name=collection_name,
    vectors_config={
        "record": rest.VectorParams(
            distance=rest.Distance.COSINE,
            size=vector_size,
        )
    },
)

# Populate collection with vectors

qdrant.upsert(
    collection_name=collection_name,
    points=[
        rest.PointStruct(
            id=k,
            vector={
                "record": v["embedding"],
            },
            payload=v.to_dict(),
        )
        for k, v in article_df.iterrows()
    ],
)
